from .base import PoseEstimator, center_ellipse
import numpy as np
from pose import initial_pose


class ICP(PoseEstimator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __name__(self):
        return "icp"

    def forward(self, C_3d, C_2d, *args,**ignore_kwargs):
        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        retval, R, T = self.estimator(X_cnt, x_cnt)
        if not retval:
            return False, None, None
        # 批量求解新的陨石坑中心
        C_2d = np.array(C_2d)
        rc = -np.linalg.inv(np.matmul(np.matmul(self.K.T, C_2d), self.K)) @ R[:, 2]
        d = -rc @ R[:, 2] / (R[:, 2] @ T)
        X_cnt_new = rc / d[:, None]
        return True, *self.solve(X_cnt, X_cnt_new)

    def solve(self, X, Y):
        """
        Copy from https://zhuanlan.zhihu.com/p/111322916
        Solve Y = RX + t

        Input:
            X: Nx3 numpy array of N points in camera coordinate
            Y: Nx3 numpy array of N points in world coordinate
        Returns:
            R: 3x3 numpy array describing camera orientation in the world (R_wc)
            t: 1x3 numpy array describing camera translation in the world (t_wc)
        """
        cx, cy = X.sum(axis=0) / Y.shape[0], Y.mean(axis=0)
        x, y = np.subtract(X, cx), np.subtract(Y, cy)
        w = x.transpose() @ y
        u, s, vh = np.linalg.svd(w)
        # 对角矩阵修正反射
        ide = np.eye(3)
        ide[2][2] = np.linalg.det(vh.transpose() @ u.transpose())
        # 求解旋转矩阵和平移向量
        R = vh.transpose() @ ide @ u.transpose()
        t = cy - R @ cx
        return R, t

    def estimator(self, database, measurements):
        """
        K: 相机内参
        database: 三维点库
        measurements: 观测到的二维点
        """
        assert len(database) == len(measurements)
        assert measurements.shape[1] == 2
        assert database.shape[1] == 3
        # 求解相机位姿
        P_ba = initial_pose(self.K, measurements, database)
        if P_ba[0] is None:
            return False, None, None
        return True, *P_ba
